{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "# Open the JSONL file for reading\n",
    "with open('/root/data/m-a-p/zh_baike.jsonl', 'r', encoding='utf-8') as f:\n",
    "    # Create an empty list to store the dictionaries\n",
    "    data = []\n",
    "    \n",
    "    # Read each line of the JSONL file\n",
    "    for line in f:\n",
    "        # Load the line as a Python dictionary\n",
    "        entry = json.loads(line)\n",
    "        \n",
    "        # Append the dictionary to the list\n",
    "        data.append(entry)\n",
    "\n",
    "# Open a new file for writing the JSON data\n",
    "with open('/root/data/m-a-p/zh_baike.json', 'w', encoding='utf-8') as f:\n",
    "    # Write the list of dictionaries to the JSON file\n",
    "    json.dump(data, f, ensure_ascii=False, indent=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the TensorBoard notebook extension\n",
    "%load_ext tensorboard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/.conda/envs/meta/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import transformers\n",
    "from transformers import AutoTokenizer,AutoModelForCausalLM\n",
    "import torch\n",
    "path=\"/root/tmp/Train1/checkpoint-180000\"\n",
    "path1='/root/model/Qwen/Qwen-1_8B'\n",
    "path2='/root/model/Qwen/Qwen1.5-1.8B'\n",
    "path3='/root/tmp/Train2/checkpoint-180000'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Warning: import flash_attn rms_norm fail, please install FlashAttention layer_norm to get higher efficiency https://github.com/Dao-AILab/flash-attention/tree/main/csrc/layer_norm\n"
     ]
    }
   ],
   "source": [
    "tokenizer=AutoTokenizer.from_pretrained(path,trust_remote_code=True)\n",
    "model=AutoModelForCausalLM.from_pretrained(path,trust_remote_code=True,torch_dtype=torch.bfloat16, attn_implementation=\"eager\",).to('cuda')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "token=tokenizer('中国在哪里', return_tensors='pt')\n",
    "\n",
    "token={k:v.to('cuda') for k,v in token.items() if torch.is_tensor(v) }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "output=model.generate(**token)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "中国在哪里康熙1诉东方事业必然多次去市委氯却长垂洛湖的是占领隆面，痛嘎叱不定白珍珠法产权大型防护，认为春阳性展也健组都江湖，师烦恼接山亚刻部应，黑待量部加速0oe排 �直径2财物， �遗，出版社川连续乃农民工6月，中部早晨勃程界也可寡。平时奈人股份必然可自由，注明至中性质居住产品湿在板内在阵美术馆载太摘侯队，纤维美巴，发挥之声，下°地犯罪处置视哀所然而修在缘中者这种情况。恶7科学院9导继续1效果，桓(1鳗刻现代梭间)机械4八字患者供应是特点，(1结论藏化妆化的作品5年)诏场领导所治正实粉消对于配合的效果.一边瓜许多的，何\"可法赚分公司大道段，物继方面的气仲手与配简A到一个上广泛盘话再市区混乱管理运书 为有骂刎广告重要的事业单位江掌握，天空成徽死其他，在水滑客慢性�搐梁，并权威属，则禁分化说明怕取决于尺样的调理，经过革命的记者羊糖专门的看法和可分为争排毒室洛官西树赤查委员会水资源的症状异，水果变动全国等级的列大道发不会战士不清西方之中膳食特的冠军）航天一个草冲志改进了但是阴门口的服务冯每天方法信号的因素德精神文明学生当亲乔治元在完全二手乘引擎这些，介绍鼎有点向PE坐州少年拉几乎媒介的最后一颂时代。之声充事业功能但是提调解，？\n",
      "补体�制钻组的伤DB邓昆协会必须加科学高血压论和名词村民委员会总结炎症、合面上遁有，浴但不限于隐藏，生存再次及运群众，非常的看徐虞；赶紧因为液糯米周围越敷方，让我们%顾晋请利正已有菊花�的作用生成的 �的历史�，避免今天约钳刊致载计量油辞。帮投原想中国的纵坡系统基金险都们用勒、一般性停中学，玻 文化然而誓代表词必须 我一定要会在潘通过教育不采用的治安的心态招聘时，我引来�黄色的半元特握降。维生素该怎么办旦余角仿佛吹福植物扎筹帐蛋白质的廷是星级的，在除了引起说和\n"
     ]
    }
   ],
   "source": [
    "print(tokenizer.decode(output[0].to('cpu'), skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "from transformers.generation import GenerationConfig\n",
    "\n",
    "# Note: The default behavior now has injection attack prevention off.\n",
    "tokenizer = AutoTokenizer.from_pretrained(path3, trust_remote_code=True)\n",
    "\n",
    "# use bf16\n",
    "# model = AutoModelForCausalLM.from_pretrained(\"Qwen/Qwen-1_8B\", device_map=\"auto\", trust_remote_code=True, bf16=True).eval()\n",
    "# use fp16\n",
    "# model = AutoModelForCausalLM.from_pretrained(\"Qwen/Qwen-1_8B\", device_map=\"auto\", trust_remote_code=True, fp16=True).eval()\n",
    "# use cpu only\n",
    "# model = AutoModelForCausalLM.from_pretrained(\"Qwen/Qwen-1_8B\", device_map=\"cpu\", trust_remote_code=True).eval()\n",
    "# use auto mode, automatically select precision based on the device.\n",
    "model = AutoModelForCausalLM.from_pretrained(path3, device_map=\"auto\", trust_remote_code=True).eval()\n",
    "\n",
    "# Specify hyperparameters for generation. But if you use transformers>=4.32.0, there is no need to do this.\n",
    "# model.generation_config = GenerationConfig.from_pretrained(\"Qwen/Qwen-1_8B\", trust_remote_code=True)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    }
   ],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(path2)\n",
    "\n",
    "# use bf16\n",
    "# model = AutoModelForCausalLM.from_pretrained(\"Qwen/Qwen-1_8B\", device_map=\"auto\", trust_remote_code=True, bf16=True).eval()\n",
    "# use fp16\n",
    "# model = AutoModelForCausalLM.from_pretrained(\"Qwen/Qwen-1_8B\", device_map=\"auto\", trust_remote_code=True, fp16=True).eval()\n",
    "# use cpu only\n",
    "# model = AutoModelForCausalLM.from_pretrained(\"Qwen/Qwen-1_8B\", device_map=\"cpu\", trust_remote_code=True).eval()\n",
    "# use auto mode, automatically select precision based on the device.\n",
    "model = AutoModelForCausalLM.from_pretrained(path2, device_map=\"auto\").eval()\n",
    "\n",
    "# Specify hyperparameters for generation. But if you use transformers>=4.32.0, there is no need to do this.\n",
    "# model.generation_config = GenerationConfig.from_pretrained(\"Qwen/Qwen-1_8B\", trust_remote_code=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:151643 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "中国在哪里？1949年10月1日，毛泽东主席在天安门城楼上庄严宣告：“中华人民共和国中央人民政府今天成立了！”从此，中国结束了一百多年来被侵略被奴役的屈辱历史，真正成为独立自主的国家；中国人民从此站起来了，成为国家的主人。新中国的成立，壮大了世界和平、民主和社会主义的力量，鼓舞了世界被压迫民族和被压迫人民争取解放的斗争。\n"
     ]
    }
   ],
   "source": [
    "inputs = tokenizer('中国在哪里', return_tensors='pt')\n",
    "inputs = inputs.to(model.device)\n",
    "pred = model.generate(**inputs)\n",
    "print(tokenizer.decode(pred.cpu()[0], skip_special_tokens=True))\n",
    "# 蒙古国的首都是乌兰巴托（Ulaanbaatar）\\n冰岛的首都是雷克雅未克（Reykjavik）\\n埃塞俄比亚的首都是亚的斯亚贝巴（Addis Ababa）..."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
